from pathlib import Path
from abc import ABC, abstractmethod
import numpy as np
import rclpy
from rclpy.node import Node
from rclpy.qos import QoSProfile
import rclpy.qos
import message_filters
from sensor_msgs.msg import Image, CameraInfo, PointCloud2
COLOR_TOPICS = ["/camera1/color/image_raw", "/camera2/color/image_raw", "/camera3/color/image_raw", "/camera4/color/image_raw"]
DEPTH_TOPICS = ["/camera1/depth/image_raw", "/camera2/depth/image_raw", "/camera3/depth/image_raw", "/camera4/depth/image_raw"]
POINTS_TOPICS = ["/camera1/depth/points", "/camera2/depth/points", "/camera3/depth/points", "/camera4/depth/points"]
COLOR_INFO_TOPICS = ["/camera1/color/camera_info", "/camera2/color/camera_info", "/camera3/color/camera_info", "/camera4/color/camera_info"]
DEPTH_INFO_TOPICS = ["/camera1/depth/camera_info", "/camera2/depth/camera_info", "/camera3/depth/camera_info", "/camera4/depth/camera_info"]


class MultiTopicSynchronizer(Node, ABC):
    """
    Generic class that implements time synchronization for multiple ROS topic messages.
    """
    def __init__(self, node_name, namespace=''):
        super().__init__(node_name=node_name, namespace=namespace)

    def init(self, topics_and_msgtypes, queue_size=10, slop=None):
        """
        Initialize the MultiTopicSynchronizer.

        Parameters:
            topics_and_msgtypes (list[tuple[str, msg_type]]): A list of topics to be synchronized along with their corresponding message types.
            queue_size (int): The size of the synchronizer message queue.
            slop (float): The allowable timestamp deviation in seconds.If None, use ExactTimeSynchronizer instead of ApproximateTimeSynchronizer.
        """
        # Create a list of subscribers
        self.subscribers = []
        for topic, message_type in topics_and_msgtypes:
            qos_profile = QoSProfile(depth=queue_size, history=rclpy.qos.HistoryPolicy.KEEP_LAST,
                                     durability=rclpy.qos.DurabilityPolicy.VOLATILE,
                                     reliability=rclpy.qos.ReliabilityPolicy.BEST_EFFORT)
            subscriber = message_filters.Subscriber(self, message_type, topic, qos_profile=qos_profile)
            self.subscribers.append(subscriber)

        if slop is None:
            self.synchronizer = message_filters.TimeSynchronizer(self.subscribers, queue_size)
        else:
            self.synchronizer = message_filters.ApproximateTimeSynchronizer(self.subscribers, queue_size, slop)
        return self

    def start(self):
        """
        Start the MultiTopicSynchronizer.
        """
        self.synchronizer.registerCallback(self.sync_callback)

    @abstractmethod
    def sync_callback(self, *msgs):
        """
        Synchronized message callback function.

        Parameters:
            *msgs: A list of synchronized messages in the order of the topics list.
        """
        # Process the synchronized messages here.
        pass


class ColorDepthSyncer(MultiTopicSynchronizer):
    def __init__(self, node_name='pair_syncer', namespace=''):
        super().__init__(node_name=node_name, namespace=namespace)
        self.declare_parameter("use_point_cloud", False)
        self.declare_parameter("color_topics", COLOR_TOPICS)
        self.declare_parameter("depth_topics", DEPTH_TOPICS)
        self.declare_parameter("points_topics", POINTS_TOPICS)
        self.declare_parameter("color_info_topics", COLOR_INFO_TOPICS)
        self.declare_parameter("depth_info_topics", DEPTH_INFO_TOPICS)
        self.num_color = 0
        self.num_camerainfo = 0
        self.camera_intrinsics = None
        self.camera_info_qos_profile = QoSProfile(depth=1, history=rclpy.qos.HistoryPolicy.KEEP_LAST,
                                                  durability=rclpy.qos.DurabilityPolicy.VOLATILE,
                                                  reliability=rclpy.qos.ReliabilityPolicy.BEST_EFFORT)

    def init(self, queue_size=10, slop=None):
        self.use_point_cloud = self.get_parameter("use_point_cloud").value
        color_topics = self.get_parameter("color_topics").value
        depth_topics = self.get_parameter("depth_topics").value
        points_topics = self.get_parameter("points_topics").value
        color_info_topics = self.get_parameter("color_info_topics").value
        depth_info_topics = self.get_parameter("depth_info_topics").value
        self.num_color = len(color_topics)
        self.num_camerainfo = len(color_topics)

        if self.use_point_cloud:
            if not (self.num_color == len(color_info_topics)):
                raise ValueError("The number of color and camera info topics must be the same.")
            topics_and_msgtypes = [(topic, Image) for topic in color_topics] + [(topic, PointCloud2) for topic in depth_topics]
            self.camera_info_topics = color_info_topics
        else:
            if not ((self.num_color  == len(color_info_topics)) and (len(depth_topics) == len(depth_info_topics))):
                raise ValueError("The number of color, depth, and camera info topics must be the same.")
            topics_and_msgtypes = [(topic, Image) for topic in color_topics] + [(topic, Image) for topic in depth_topics]
            self.num_camerainfo += len(depth_info_topics)
            self.camera_info_topics = color_info_topics + depth_info_topics
        super().init(topics_and_msgtypes, queue_size=queue_size, slop=slop)
        if self.camera_intrinsics is None:
            self.camera_intrinsics = [None] * self.num_camerainfo
            self.current_info_idx = 0
            self.current_info_subscription = None
            self.subscribe_to_next_info()
        else:
            self.start()
        return self

    def subscribe_to_next_info(self):
        if self.current_info_idx < self.num_camerainfo:
            self.current_info_subscription = self.create_subscription(
                CameraInfo,
                self.camera_info_topics[self.current_info_idx],
                self.camera_info_callback,
                self.camera_info_qos_profile
            )
            self.get_logger().info(f'Waiting for camera {self.current_info_idx} info...')
        else:
            self.start()
            self.get_logger().info('All camera info received')

    def camera_info_callback(self, msg):
        self.get_logger().info("color info received")
        self.camera_intrinsics[self.current_info_idx] = np.array(msg.k).reshape(3, 3)
        self.get_logger().info("Intrinsic matrix: \n{}".format(self.camera_intrinsics[self.current_info_idx]))
        # Destroy current subscription
        self.destroy_subscription(self.current_info_subscription)
        self.current_info_subscription = None
        # Subscribe to the next camera topic
        self.current_info_idx += 1
        self.subscribe_to_next_info()

    def sync_callback(self, *msgs):
        color_msgs = msgs[:self.num_color]
        depth_msgs = msgs[self.num_color:]
        self.process_observation(color_msgs, depth_msgs)

    @abstractmethod
    def process_observation(self, color_msgs, depth_msgs):
        """
        Process the synchronized color and depth messages.

        Parameters:
            color_msgs (list[sensor_msgs.msg.Image]): A list of synchronized color messages.
            depth_msgs (list[sensor_msgs.msg.Image or sensor_msgs.msg.PointCloud2]): A list of synchronized depth messages.
        """
        pass
